package org.example;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import org.rocksdb.*;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

public class WriteRocksDBMain {

    static {
        RocksDB.loadLibrary();
    }

    private static final int BATCH_SIZE = 5000; // 批量写入更小，减少内存压力
    private static final int THREADS = Runtime.getRuntime().availableProcessors();
    private static final Gson GSON = new Gson();

    public static void main(String[] args) throws IOException, RocksDBException, InterruptedException {
        String inputFile = "md5.jsonl";  // 输入文件
        String dbPath = "rocksdb_md5";   // RocksDB 数据目录

        Options options = new Options()
                .setCreateIfMissing(true)
                .setWriteBufferSize(32 * 1024 * 1024) // 32MB 写缓冲
                .setMaxWriteBufferNumber(3)
                .setMaxBackgroundJobs(THREADS)
                .setCompressionType(CompressionType.LZ4_COMPRESSION)
                .setCompactionStyle(CompactionStyle.UNIVERSAL);

        BlockBasedTableConfig tableOptions = new BlockBasedTableConfig()
                .setBlockCache(new LRUCache(64 * 1024 * 1024)) // 64MB cache
                .setFilterPolicy(new BloomFilter(10, false));
        options.setTableFormatConfig(tableOptions);

        try (RocksDB db = RocksDB.open(options, dbPath)) {
            processFileToRocksDB(inputFile, db);
        }
        //✅ Finished writing 320,120,563 lines in 1288.33 seconds (248477.34 lines/sec) 21分钟

    }

    private static void processFileToRocksDB(String filePath, RocksDB db) throws IOException, InterruptedException {
        // 阻塞队列，最多缓存 2*THREADS 个任务，避免OOM
        ExecutorService executor = new ThreadPoolExecutor(
                THREADS, THREADS, 0L, TimeUnit.MILLISECONDS,
                new ArrayBlockingQueue<>(THREADS * 2),
                new ThreadPoolExecutor.CallerRunsPolicy() // 如果满了就在当前线程执行
        );

        AtomicLong totalLines = new AtomicLong(0);
        long startTime = System.currentTimeMillis();

        try (BufferedReader reader = new BufferedReader(new FileReader(filePath, StandardCharsets.UTF_8))) {
            List<String> batch = new ArrayList<>(BATCH_SIZE);
            String line;
            while ((line = reader.readLine()) != null) {
                batch.add(line);
                if (batch.size() >= BATCH_SIZE) {
                    List<String> batchCopy = new ArrayList<>(batch);
                    executor.submit(() -> writeBatch(batchCopy, db, totalLines, startTime));
                    batch.clear();
                }
            }
            if (!batch.isEmpty()) {
                List<String> batchCopy = new ArrayList<>(batch);
                executor.submit(() -> writeBatch(batchCopy, db, totalLines, startTime));
            }
        }

        executor.shutdown();
        executor.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);

        long totalTime = System.currentTimeMillis() - startTime;
        System.out.printf("✅ Finished writing %,d lines in %.2f seconds (%.2f lines/sec)%n",
                totalLines.get(), totalTime / 1000.0, totalLines.get() / (totalTime / 1000.0));
    }

    private static void writeBatch(List<String> batch, RocksDB db, AtomicLong totalLines, long startTime) {
        try (WriteBatch writeBatch = new WriteBatch()) {
            for (String line : batch) {
                try {
                    JsonObject obj = GSON.fromJson(line, JsonObject.class);
                    String md5 = obj.get("md5").getAsString();
                    String text = obj.get("text").getAsString();
                    writeBatch.put(md5.getBytes(StandardCharsets.UTF_8), text.getBytes(StandardCharsets.UTF_8));
                } catch (Exception e) {
                    System.err.println("⚠️ Failed to parse line: " + e.getMessage());
                }
            }
            // 用全局 WriteOptions，禁用 WAL 提升速度
            try (WriteOptions writeOptions = new WriteOptions().setDisableWAL(true)) {
                db.write(writeOptions, writeBatch);
            }

            long count = totalLines.addAndGet(batch.size());
            if (count % 1_000_000 < BATCH_SIZE) {
                long elapsed = System.currentTimeMillis() - startTime;
                double speed = count / (elapsed / 1000.0);
                System.out.printf("📊 Written %,d lines, speed: %.2f lines/sec%n", count, speed);
            }
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }
}
